#include "TestEigens.h"
#include "../../OclLapack/OclLapack.h"
#include <vector>
#include <float.h>
#include <string.h>

#include "../common/TestUtils.h"

using namespace std;

const int RANDOM_TEST_SIZE = 100;
const int NUM_RANDOM_REPEATS = 20;

Lapack *TestEigens::lapack = NULL;
OclLapack *TestEigens::lapackCPU = NULL;
OclLapack *TestEigens::lapackGPU = NULL;

#define CHECK_LAPACK(info) CPPUNIT_ASSERT(info == 0)

CPPUNIT_TEST_SUITE_REGISTRATION(TestEigens);

Lapack* TestEigens::getLapack() {
	if (lapack == NULL) {
		lapack = new Lapack();
	}
	return lapack;
}

Lapack* TestEigens::getOclLapackCPU() {
	if (lapackCPU == NULL) {
		vector<cl_device_id> devices;
		enumOpenCLDevices(devices, CL_DEVICE_TYPE_CPU);
		lapackCPU = new OclLapack(devices[1]);
	}
	return lapackCPU;
}

Lapack* TestEigens::getOclLapackGPU() {
	if (lapackGPU == NULL) {
		vector<cl_device_id> devices;
		enumOpenCLDevices(devices, CL_DEVICE_TYPE_GPU);
		lapackGPU = new OclLapack(devices[0]);
		//lapackGPU = new OclLapack(CL_DEVICE_TYPE_GPU);
	}
	return lapackGPU;
}

TestEigens::TestEigens() { 
}

TestEigens::~TestEigens() {
}

void TestEigens::setUp() {
}

void TestEigens::tearDown() {
}

void printMatrix(int height, int width, float* matrix) {
	for (int i = 0; i < height; i++) {
		for (int j = 0; j < width; j++) {
			printf("%-9.6f ", matrix[i + j * height]);
		}
		printf("\n");
	}
}

void printVector(int n, float *v) {
	for (int i = 0; i < n; i++) {
		printf("%-9.6f ", v[i]);
	}
	printf("\n");
}

void testEigens(Lapack &lapack, int size, float* matrix, float* outEigenvalues, float* outEigenvectors, int* outNumEigens) {
		
	vector<float> backupMatrix(size, size);
	backupMatrix.assign(&matrix[0], &matrix[size * size]);
	
	try {
			
		// determine work size
		int lwork = -1;
		int info = 0;
		float workLength;
		int tempInt = 1;
		float tempFloat;
		lapack.ssyevx("Vectors", "All", "L", &size, &matrix[0], &size,
				NULL, NULL, NULL, NULL, NULL, NULL, &tempFloat, &tempFloat, &size, 
				&workLength, &lwork, NULL, &tempInt, &info);
		lwork = (int)workLength;
			
		// calculate
		int numEigens;
		vector<float> eigenvectors(size * size);
		vector<float> eigens(size);
		vector<int> ifail(size);
		vector<float> work(lwork);
		int intWorkSize = size * 5;
		vector<int> iwork(intWorkSize);
		float tolerance = FLT_EPSILON;
				
		lapack.ssyevx("Vectors", "All", "L", &size, &matrix[0], &size,
				NULL, NULL, NULL, NULL, &tolerance,
				&numEigens,  &eigens[0], &eigenvectors[0], &size, 
				&work[0], &lwork, &iwork[0], &ifail[0], &info);
		CHECK_LAPACK(info);
		
		if (outEigenvalues != NULL) {
			memcpy(outEigenvalues, &eigens[0], numEigens * sizeof(float));
		}
		
		if (outEigenvectors != NULL) {
			memcpy(outEigenvectors, &eigenvectors[0], numEigens * size * sizeof(float));
		}
		
		if (outNumEigens != NULL) {
			*outNumEigens = numEigens;
		}
				
		// test eigenvalue/eigenvector properties
		vector<float> temp(size);
		memcpy(matrix, &backupMatrix[0], size * size * sizeof(float));
		
		for (int i = 0; i < numEigens; i++) {									
			int one = 1;
			float minus1 = -1;
			
			// temp = eigenvectors[i]			
			vector<float>::iterator from = eigenvectors.begin() + i * size;
			temp.assign(from, from + size);
						
			// temp = -matrix * eigenvectors[i] + eigens[i] * temp
			lapack.ssymv("L", &size, &minus1, 
					&matrix[0], &size,
					&eigenvectors[i * size], &one, 
					&eigens[i],
					&temp[0], &one);
			
			float norm = lapack.snrm2(&size, &temp[0], &one);
			CPPUNIT_ASSERT(compare(norm, 0, 2E-4f) == 0);
		}
		
		// test eigenvector orthogonality
		for (int i = 0; i < numEigens; i++) {
			for (int j = i + 1; j < numEigens; j++) {
				int one = 1;
				float dot = lapack.sdot(&size, &eigenvectors[i * size], &one,
						&eigenvectors[j * size], &one);
				CPPUNIT_ASSERT(compare(dot, 0, 5E-6f) == 0);
			}
		}		
		
	} catch (Exception& e) {
		CPPUNIT_ASSERT_MESSAGE(e.what(), false);
	}
}

void checkKnownEigens(Lapack &lapack) {
	float matrix[] = {
		5, 4, 3, 2, 1,
		4, 5, 4, 3, 2,
		3, 4, 5, 4, 3,
		2, 3, 4, 5, 4,
		1, 2, 3, 4, 5
	};
		
	float correctEigenvalues[] = { 0.548432, 0.763932, 1.273771, 5.236068, 17.177798 };
		
	int size = sizeof(correctEigenvalues) / sizeof(correctEigenvalues[0]);
		
	try {		
		vector<float> eigenvalues(size);
		int numEigens;
		testEigens(lapack, size, matrix, &eigenvalues[0], NULL, &numEigens);
				
		CPPUNIT_ASSERT_EQUAL(size, numEigens);
		CPPUNIT_ASSERT(compare(&eigenvalues[0], correctEigenvalues, eigenvalues.size(), 20 * FLT_EPSILON) == 0);
	} catch (Exception& e) {
		CPPUNIT_ASSERT_MESSAGE(e.what(), false);
	}
}

void testRandomMatrix(Lapack& lapack, int size) {
	// fill with random values
	vector<float> matrix(size * size);
	for (int i = 0; i < size; i++) {
		matrix[i * size + i] = 1;
		for (int j = i + 1; j < size; j++) {
			matrix[i * size + j] = matrix[j * size + i] = getRandom();
		}
	}
	
	testEigens(lapack, size, &matrix[0], NULL, NULL, NULL);
}

void testIntervalOfEigens(Lapack &lapack) {
	float correctEigenvalues[] = { 0.548432, 0.763932, 1.273771, 5.236068, 17.177798 };
	
	float d[] = { 5, 13.0666656, 5.30520487, 0.968692899, 0.659436226 };
	float e[] = { -5.4772253, 4.4191494, 0.567749798, -0.143790632 };
	int size = sizeof(d) / sizeof(d[0]);
	
	int lowIndex = 2;
	int upIndex = 4;
	float tolerance = FLT_EPSILON;
	int numEigens;
	int numSplit;
	vector<float> eigenvalues(size);
	vector<int> iblock(size), isplit(size);
	vector<float> work(4 * size);
	vector<int> iwork(3 * size);
	int info;
	
	lapack.sstebz("Index", "Entire matrix", &size,
			NULL, NULL,
			&lowIndex, &upIndex,
			&tolerance, 
			d, e, &numEigens,
			&numSplit,
			&eigenvalues[0], 
			&iblock[0], &isplit[0], 
			&work[0], &iwork[0], &info);
	
	CHECK_LAPACK(info);
	CPPUNIT_ASSERT(compare(&eigenvalues[0], &correctEigenvalues[lowIndex - 1], upIndex - lowIndex + 1, 10 * FLT_EPSILON) == 0);
}

void TestEigens::testKnownEigensLapack() {
	checkKnownEigens(*getLapack());	
}

void TestEigens::testKnownEigensCPU() {
	Lapack *l = getOclLapackCPU();
	checkKnownEigens(*l);		
}

void TestEigens::testKnownEigensGPU() {
	checkKnownEigens(*getOclLapackGPU());	
}

void TestEigens::testRandomEigensLapack() {
	srand(102452);	
	for (int i = 0; i < NUM_RANDOM_REPEATS; i++) {
		testRandomMatrix(*getLapack(), RANDOM_TEST_SIZE);
	}
}

void TestEigens::testRandomEigensCPU() {
	srand(102452);
	int size = 200;	
	for (int i = 0; i < NUM_RANDOM_REPEATS; i++) {
		testRandomMatrix(*getOclLapackCPU(), RANDOM_TEST_SIZE);
	}
}

void TestEigens::testRandomEigensGPU() {
	srand(102452);
	int size = 200;	
	for (int i = 0; i < NUM_RANDOM_REPEATS; i++) {
		testRandomMatrix(*getOclLapackGPU(), RANDOM_TEST_SIZE);
	}
}

void TestEigens::testIntervalOfEigensLapack() {	
	testIntervalOfEigens(*getLapack());	
}

void TestEigens::testIntervalOfEigensCPU() {
	testIntervalOfEigens(*getOclLapackCPU());	
}

void TestEigens::testIntervalOfEigensGPU() {
	testIntervalOfEigens(*getOclLapackGPU());	
}

void TestEigens::compareSstebz(int size) {
	int numEigens, numSplits;
	int info;
	float eps = FLT_EPSILON;
		
	vector<float> d(size);
	vector<float> e(size);	
	vector<int> iblock(size), isplit(size);
	
	vector<float> work(size * 4);
	vector<int> iwork(size * 3); 
	
	vector<float> eigenvaluesLapack(size);
	vector<float> eigenvaluesCPU(size);
	vector<float> eigenvaluesGPU(size);
	
	for (int i = 0; i < size; i++) {
		d[i] = getRandom();
		e[i] = getRandom();
	}
	
	char lapack[32], oclCpu[32], oclGpu[32];
	sprintf(lapack, "test-sstebz-lapack-%d", size);
	sprintf(oclCpu, "test-sstebz-ocl-cpu-%d", size);
	sprintf(oclGpu, "test-sstebz-ocl-gpu-%d", size);
	
	{	
		Measure m(lapack);
		getLapack()->sstebz("A", "E", &size, 
			NULL, NULL, NULL, NULL, 
			&eps, &d[0], &e[0], &numEigens,
			&numSplits, &eigenvaluesLapack[0], 
			&iblock[0], &isplit[0], &work[0], &iwork[0], &info);
	}
	
//	Lapack *lapackCPU = getOclLapackCPU();
//	{
//		Measure m(oclCpu);
//		lapackCPU->sstebz("A", "E", &size, 
//			NULL, NULL, NULL, NULL, 
//			&eps, &d[0], &e[0], &numEigens,
//			&numSplits, &eigenvaluesCPU[0], 
//			&iblock[0], &isplit[0], &work[0], &iwork[0], &info);
//	}
	
	Lapack *lapackGPU = getOclLapackGPU();
	{
		Measure m(oclGpu);
		lapackGPU->sstebz("A", "E", &size, 
			NULL, NULL, NULL, NULL, 
			&eps, &d[0], &e[0], &numEigens,
			&numSplits, &eigenvaluesGPU[0], 
			&iblock[0], &isplit[0], &work[0], &iwork[0], &info);
	}
	
	//CPPUNIT_ASSERT(compare(&eigenvaluesLapack[0], &eigenvaluesCPU[0], numEigens, 1E-5) == 0);
	//CPPUNIT_ASSERT(compare(&eigenvaluesLapack[0], &eigenvaluesGPU[0], numEigens, 1E-5) == 0);
}

void TestEigens::compareSstebz() {
	srand(1384562);
	for (int size = 64; size <= 4096; size *= 2) {
		for (int i = 0; i < 50; i++) {
			compareSstebz(size);
		}
	}
}

void checkKnownSstebzOnly(Lapack &lapack) {
	float correctEigenvalues[] = { 0.548432, 0.763932, 1.273771, 5.236068, 17.177798 };
	
	float d[] = { 5, 13.0666656, 5.30520487, 0.968692899, 0.659436226 };
	float e[] = { -5.4772253, 4.4191494, 0.567749798, -0.143790632 };
	int size = sizeof(d) / sizeof(d[0]);
	
	
	float tolerance = FLT_EPSILON;
	int numEigens;
	int numSplit;
	vector<float> eigenvalues(size);
	vector<int> iblock(size), isplit(size);
	vector<float> work(4 * size);
	vector<int> iwork(3 * size);
	int info;
	
	lapack.sstebz("All", "B", &size,
			NULL, NULL, NULL, NULL,
			&tolerance, 
			d, e, &numEigens,
			&numSplit,
			&eigenvalues[0], 
			&iblock[0], &isplit[0], 
			&work[0], &iwork[0], &info);
	
	CHECK_LAPACK(info);
	CPPUNIT_ASSERT(compare(&eigenvalues[0], &correctEigenvalues[0], size, 5E-6) == 0);
}

void TestEigens::testKnownSstebzOnlyOnGPU() {
	checkKnownSstebzOnly(*getOclLapackGPU());
}

